import pandas as pd
import clean_data
import numpy as np
from sklearn.model_selection import StratifiedKFold  # 分层K折交叉验证类似于网格搜索cv=折数
import joblib
from sklearn.model_selection import GridSearchCV
import xgboost as xgb
from sklearn.preprocessing import LabelEncoder  # 文字分类
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import transform_train

# 数据读取
x_xgb, y = transform_train.train_do()
x_xgb = x_xgb.iloc[:, 0:20]
def xgb_train():
    """
    参数解释
        n_estimators:决策树数量
        max_depth:决策树分支 
        learning_rate:学习率
        objective='binary:logistic':二分类
        random_state:随机种子
    
    """
    # 模型训练
    # 交叉验证+网格搜索
    model_xgb = xgb.XGBClassifier(objective='binary:logistic')
    param_grid = {
        'n_estimators': [150, 200, 250],
        'max_depth': [3, 5],
        'learning_rate': [0.05, 0.06, 0.08],
        'random_state': [22]
    }
    grid_search = GridSearchCV(estimator=model_xgb, param_grid=param_grid, cv=5, n_jobs=-1)
    grid_search.fit(x_xgb, y)
    print(grid_search.best_params_)
    joblib.dump(grid_search.best_estimator_, '../model/xgb.pkl')

if __name__ == '__main__':
    xgb_train()
